import numpy, h5py
import pycuda.driver as cuda
import pycuda.autoinit

from spa import ManyRotate2D
from PIL import Image

N,K = 128,5
nrot = 24

hfile = h5py.File('hdf5/image_data.hdf5')
image = hfile['Images']['sar00001'][:][1:N+1,1:N+1]
hfile.close()
image = image.copy().astype(numpy.float32)
image -= image.mean()
image /= numpy.sqrt(image.var())

min_val,scale = -2.0, 255.0/4.0

I = Image.new("L",(N,N))
for i in range(N):
	for j in range(N):
		val = scale*(image[i,j] - min_val)
		val = max(0.0,min(255.0,val))
		I.putpixel((j,i),val)
I.save('pngmrot/imag1.png')
		

image_gpu = cuda.mem_alloc(image.nbytes)
cuda.memcpy_htod(image_gpu,image)

rimages = numpy.zeros((nrot,N,N),numpy.float32)
rimages_gpu = cuda.mem_alloc(rimages.nbytes)

angle = 2.0*numpy.pi/float(nrot)

ManyRotate2D(image_gpu, rimages_gpu, numpy.float32(angle),
	numpy.int32(N), numpy.int32(K), numpy.float32(0.975),
	block=(N,1,1), grid=(N,nrot))

cuda.memcpy_dtoh(rimages, rimages_gpu)


for k in range(nrot):
	I = Image.new("L",(N,N))
	for i in range(N):
		for j in range(N):
			val = scale*(rimages[k,i,j] - min_val)
			val = max(0.0,min(255.0,val))
			I.putpixel((j,i),val)
	I.save('pngmrot/rotation' + ('%03d' % (k+1,)) + '.png')
